{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "private_outputs": true,
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "gpuClass": "standard"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Stable Diffusion Textual Inversion - Concept Library navigation and usage\n",
        "\n",
        "Navigate through the [public library of concepts](https://huggingface.co/sd-concepts-library) and use Stable Diffusion with custom concepts. 🤗 Hugging Face [🧨 Diffusers library](https://github.com/huggingface/diffusers). \n",
        "\n",
        "![Textual Inversion example](https://textual-inversion.github.io/static/images/editing/colorful_teapot.JPG)\n",
        "_By using just 3-5 images new concepts can be taught to Stable Diffusion and the model personalized on your own images_ \n",
        "\n",
        "If you would like to teach Stable Diffusion your own concepts, check out the [training notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_textual_inversion_training.ipynb)\n"
      ],
      "metadata": {
        "id": "forgOfmQeA-l"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Initial setup"
      ],
      "metadata": {
        "id": "CGO3td9-LZzY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install the required libs\n",
        "!pip install -qq diffusers==0.4.1 transformers ftfy gradio wget"
      ],
      "metadata": {
        "id": "FQOlXb7Pdbj2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Login to the Hugging Face Hub\n",
        "#@markdown If you haven't yet, [you have to first acknowledge and agree to the model LICENSE before using it](https://huggingface.co/CompVis/stable-diffusion-v1-4) \n",
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "rnhKBvKidtxK"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Prepare the Concepts Library to be used\n",
        "\n",
        "import requests\n",
        "import os\n",
        "import gradio as gr\n",
        "import wget\n",
        "import torch\n",
        "from diffusers import StableDiffusionPipeline\n",
        "from huggingface_hub import HfApi\n",
        "from transformers import CLIPTextModel, CLIPTokenizer\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "api = HfApi()\n",
        "models_list = api.list_models(author=\"sd-concepts-library\")\n",
        "models = []\n",
        "\n",
        "pipe = StableDiffusionPipeline.from_pretrained(\"CompVis/stable-diffusion-v1-4\", revision=\"fp16\", torch_dtype=torch.float16).to(\"cuda\")\n",
        "\n",
        "def load_learned_embed_in_clip(learned_embeds_path, text_encoder, tokenizer, token=None):\n",
        "  loaded_learned_embeds = torch.load(learned_embeds_path, map_location=\"cpu\")\n",
        "  \n",
        "  # separate token and the embeds\n",
        "  trained_token = list(loaded_learned_embeds.keys())[0]\n",
        "  embeds = loaded_learned_embeds[trained_token]\n",
        "\n",
        "  # cast to dtype of text_encoder\n",
        "  dtype = text_encoder.get_input_embeddings().weight.dtype\n",
        "  embeds.to(dtype)\n",
        "\n",
        "  # add the token in tokenizer\n",
        "  token = token if token is not None else trained_token\n",
        "  num_added_tokens = tokenizer.add_tokens(token)\n",
        "  i = 1\n",
        "  while(num_added_tokens == 0):\n",
        "    print(f\"The tokenizer already contains the token {token}.\")\n",
        "    token = f\"{token[:-1]}-{i}>\"\n",
        "    print(f\"Attempting to add the token {token}.\")\n",
        "    num_added_tokens = tokenizer.add_tokens(token)\n",
        "    i+=1\n",
        "  \n",
        "  # resize the token embeddings\n",
        "  text_encoder.resize_token_embeddings(len(tokenizer))\n",
        "  \n",
        "  # get the id for the token and assign the embeds\n",
        "  token_id = tokenizer.convert_tokens_to_ids(token)\n",
        "  text_encoder.get_input_embeddings().weight.data[token_id] = embeds\n",
        "  return token\n",
        "\n",
        "print(\"Setting up the public library\")\n",
        "for model in tqdm(models_list):\n",
        "  model_content = {}\n",
        "  model_id = model.modelId\n",
        "  model_content[\"id\"] = model_id\n",
        "  embeds_url = f\"https://huggingface.co/{model_id}/resolve/main/learned_embeds.bin\"\n",
        "  os.makedirs(model_id,exist_ok = True)\n",
        "  if not os.path.exists(f\"{model_id}/learned_embeds.bin\"):\n",
        "    try:\n",
        "      wget.download(embeds_url, out=model_id)\n",
        "    except:\n",
        "      continue\n",
        "  token_identifier = f\"https://huggingface.co/{model_id}/raw/main/token_identifier.txt\"\n",
        "  response = requests.get(token_identifier)\n",
        "  token_name = response.text\n",
        "  \n",
        "  concept_type = f\"https://huggingface.co/{model_id}/raw/main/type_of_concept.txt\"\n",
        "  response = requests.get(concept_type)\n",
        "  concept_name = response.text\n",
        "  model_content[\"concept_type\"] = concept_name\n",
        "  images = []\n",
        "  for i in range(4):\n",
        "    url = f\"https://huggingface.co/{model_id}/resolve/main/concept_images/{i}.jpeg\"\n",
        "    image_download = requests.get(url)\n",
        "    url_code = image_download.status_code\n",
        "    if(url_code == 200):\n",
        "      file = open(f\"{model_id}/{i}.jpeg\", \"wb\") ## Creates the file for image\n",
        "      file.write(image_download.content) ## Saves file content\n",
        "      file.close()\n",
        "      images.append(f\"{model_id}/{i}.jpeg\")\n",
        "  model_content[\"images\"] = images\n",
        "\n",
        "  learned_token = load_learned_embed_in_clip(f\"{model_id}/learned_embeds.bin\", pipe.text_encoder, pipe.tokenizer, token_name)\n",
        "  model_content[\"token\"] = learned_token\n",
        "  models.append(model_content)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "uvBezVlgfUuL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Go!\n"
      ],
      "metadata": {
        "id": "MOL_FclPLcJw"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wDG37T4ic-KK"
      },
      "outputs": [],
      "source": [
        "#@title Run the app to navigate around [the Library](https://huggingface.co/sd-concepts-library)\n",
        "#@markdown Click the `Running on public URL:` result to run the Gradio app\n",
        "\n",
        "SELECT_LABEL = \"Select concept\"\n",
        "\n",
        "def title_block(title, id):\n",
        "  return gr.Markdown(f\"### [`{title}`](https://huggingface.co/{id})\")\n",
        "\n",
        "def image_block(image_list, concept_type):\n",
        "  return gr.Gallery(\n",
        "          label=concept_type, value=image_list, elem_id=\"gallery\"\n",
        "          ).style(grid=[2], height=\"auto\")\n",
        "\n",
        "def checkbox_block():\n",
        "  checkbox = gr.Checkbox(label=SELECT_LABEL).style(container=False)\n",
        "  return checkbox\n",
        "\n",
        "def infer(text):\n",
        "  images_list = pipe(\n",
        "      text,\n",
        "      num_images_per_prompt=2,\n",
        "      num_inference_steps=50,\n",
        "      guidance_scale=7.5\n",
        "  )\n",
        "  output_images = []\n",
        "  for i, image in enumerate(images_list[\"sample\"]):\n",
        "    output_images.append(image)\n",
        "  return output_images\n",
        "  \n",
        "css = '''\n",
        ".gradio-container {font-family: 'IBM Plex Sans', sans-serif}\n",
        "#top_title{margin-bottom: .5em}\n",
        "#top_title h2{margin-bottom: 0; text-align: center}\n",
        "#main_row{flex-wrap: wrap; gap: 1em; max-height: calc(100vh - 16em); overflow-y: scroll; flex-direction: row}\n",
        "@media (min-width: 768px){#main_row > div{flex: 1 1 32%; margin-left: 0 !important}}\n",
        ".gr-prose code::before, .gr-prose code::after {content: \"\" !important}\n",
        "::-webkit-scrollbar {width: 10px}\n",
        "::-webkit-scrollbar-track {background: #f1f1f1}\n",
        "::-webkit-scrollbar-thumb {background: #888}\n",
        "::-webkit-scrollbar-thumb:hover {background: #555}\n",
        ".gr-button {white-space: nowrap}\n",
        ".gr-button:focus {\n",
        "  border-color: rgb(147 197 253 / var(--tw-border-opacity));\n",
        "  outline: none;\n",
        "  box-shadow: var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow, 0 0 #0000);\n",
        "  --tw-border-opacity: 1;\n",
        "  --tw-ring-offset-shadow: var(--tw-ring-inset) 0 0 0 var(--tw-ring-offset-width) var(--tw-ring-offset-color);\n",
        "  --tw-ring-shadow: var(--tw-ring-inset) 0 0 0 calc(3px var(--tw-ring-offset-width)) var(--tw-ring-color);\n",
        "  --tw-ring-color: rgb(191 219 254 / var(--tw-ring-opacity));\n",
        "  --tw-ring-opacity: .5;\n",
        "}\n",
        "#prompt_input{flex: 1 3 auto}\n",
        "#prompt_area{margin-bottom: .75em}\n",
        "#prompt_area > div:first-child{flex: 1 3 auto}\n",
        "'''\n",
        "examples = [\"a <cat-toy> in <madhubani-art> style\", \"a mecha robot in <line-art> style\", \"a piano being played by <bonzi>\"]\n",
        "with gr.Blocks(css=css) as demo:\n",
        "  state = gr.Variable({\n",
        "        'selected': -1\n",
        "  })\n",
        "  state = {}\n",
        "  def update_state(i):\n",
        "        global checkbox_states\n",
        "        if(checkbox_states[i]):\n",
        "          checkbox_states[i] = False\n",
        "          state[i] = False\n",
        "        else:\n",
        "          state[i] = True\n",
        "          checkbox_states[i] = True\n",
        "  gr.HTML('''\n",
        "  <div style=\"text-align: center; max-width: 720px; margin: 0 auto;\">\n",
        "              <div\n",
        "                style=\"\n",
        "                  display: inline-flex;\n",
        "                  align-items: center;\n",
        "                  gap: 0.8rem;\n",
        "                  font-size: 1.75rem;\n",
        "                \"\n",
        "              >\n",
        "                <svg\n",
        "                  width=\"0.65em\"\n",
        "                  height=\"0.65em\"\n",
        "                  viewBox=\"0 0 115 115\"\n",
        "                  fill=\"none\"\n",
        "                  xmlns=\"http://www.w3.org/2000/svg\"\n",
        "                >\n",
        "                  <rect width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect y=\"69\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"23\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"23\" y=\"69\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"46\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"46\" y=\"69\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"69\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                  <rect x=\"69\" y=\"69\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                  <rect x=\"92\" width=\"23\" height=\"23\" fill=\"#D9D9D9\"></rect>\n",
        "                  <rect x=\"92\" y=\"69\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"115\" y=\"46\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"115\" y=\"115\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"115\" y=\"69\" width=\"23\" height=\"23\" fill=\"#D9D9D9\"></rect>\n",
        "                  <rect x=\"92\" y=\"46\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"92\" y=\"115\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"92\" y=\"69\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"69\" y=\"46\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"69\" y=\"115\" width=\"23\" height=\"23\" fill=\"white\"></rect>\n",
        "                  <rect x=\"69\" y=\"69\" width=\"23\" height=\"23\" fill=\"#D9D9D9\"></rect>\n",
        "                  <rect x=\"46\" y=\"46\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                  <rect x=\"46\" y=\"115\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                  <rect x=\"46\" y=\"69\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                  <rect x=\"23\" y=\"46\" width=\"23\" height=\"23\" fill=\"#D9D9D9\"></rect>\n",
        "                  <rect x=\"23\" y=\"115\" width=\"23\" height=\"23\" fill=\"#AEAEAE\"></rect>\n",
        "                  <rect x=\"23\" y=\"69\" width=\"23\" height=\"23\" fill=\"black\"></rect>\n",
        "                </svg>\n",
        "                <h1 style=\"font-weight: 900; margin-bottom: 7px;\">\n",
        "                  Stable Diffusion Conceptualizer\n",
        "                </h1>\n",
        "              </div>\n",
        "              <p style=\"margin-bottom: 10px; font-size: 94%\">\n",
        "                Navigate through community created concepts and styles via Stable Diffusion Textual Inversion and pick yours for inference.\n",
        "                To train your own concepts and contribute to the library <a style=\"text-decoration: underline\" href=\"https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/sd_textual_inversion_training.ipynb\">check out this notebook</a>.\n",
        "              </p>\n",
        "            </div>\n",
        "  ''')\n",
        "  with gr.Row():\n",
        "        with gr.Column():\n",
        "          gr.Markdown('''\n",
        "          ### Textual-Inversion trained [concepts library](https://huggingface.co/sd-concepts-library) navigator\n",
        "          ''')\n",
        "          with gr.Row(elem_id=\"main_row\"):\n",
        "                  image_blocks = []\n",
        "                  for i, model in enumerate(models):\n",
        "                    with gr.Box().style(border=None):\n",
        "                      title_block(model[\"token\"], model[\"id\"])\n",
        "                      image_blocks.append(image_block(model[\"images\"], model[\"concept_type\"]))\n",
        "        with gr.Box():\n",
        "                with gr.Row(elem_id=\"prompt_area\").style(mobile_collapse=False, equal_height=True):\n",
        "                    text = gr.Textbox(\n",
        "                        label=\"Enter your prompt\", placeholder=\"Enter your prompt\", show_label=False, max_lines=1, elem_id=\"prompt_input\"\n",
        "                    ).style(\n",
        "                        border=(True, False, True, True),\n",
        "                        rounded=(True, False, False, True),\n",
        "                        container=False                        \n",
        "                    )\n",
        "                    btn = gr.Button(\"Run\",elem_id=\"run_btn\").style(\n",
        "                        margin=False,\n",
        "                        rounded=(False, True, True, False)\n",
        "                    )  \n",
        "                with gr.Row().style():\n",
        "                    infer_outputs = gr.Gallery(show_label=False).style(grid=[2], height=\"512px\")\n",
        "                with gr.Row():\n",
        "                  gr.HTML(\"<p style=\\\"font-size: 85%;margin-top: .75em\\\">Prompting may not work as you are used to; <code>objects</code> may need the concept added at the end.</p>\")\n",
        "                with gr.Row():\n",
        "                  gr.Examples(examples=examples, fn=infer, inputs=[text], outputs=infer_outputs, cache_examples=False)\n",
        "  checkbox_states = {}\n",
        "  inputs = [text]\n",
        "  btn.click(\n",
        "        infer,\n",
        "        inputs=inputs,\n",
        "        outputs=infer_outputs\n",
        "    )\n",
        "demo.launch(inline=False, debug=True)"
      ]
    }
  ]
}